package com.jh.hts_jdk.item;

import java.io.IOException;
import java.io.Writer;
import java.util.Arrays;
import java.util.HashMap;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;

import org.apache.commons.math3.linear.MatrixUtils;
import org.apache.commons.math3.linear.RealMatrix;

import htsjdk.samtools.util.StringUtil;

public class Matrix<R, C> {

    public String name = "";
    private RealMatrix data;
    private R[] row_name;
    private C[] col_name;
    private Lock[] row_locks;
    private Lock[] col_Locks;

    private HashMap<R, Integer> row_to_index = new HashMap<>();
    private HashMap<C, Integer> col_to_index = new HashMap<>();

    public Matrix(R[] row_name, C[] col_name, double[][] data) {
        if (row_name.length != data.length || col_name.length != data[0].length) {
            throw new IllegalArgumentException("matrix row name or col name not match data dimension");
        }
        this.data = MatrixUtils.createRealMatrix(data);
        this.row_name = row_name.clone();
        this.col_name = col_name.clone();
        indexInit();
    }

    public Matrix(R[] row_name, C[] col_name, double value) {
        this(row_name, col_name, new double[row_name.length][col_name.length]);
        for (int i = 0; i < row_name.length; i++) {
            for (int j = 0; j < col_name.length; j++) {
                this.data.setEntry(i, j, value);
            }
        }
    }

    private void indexInit() {
        row_to_index.clear();
        row_locks = new Lock[row_name.length];
        for (int i = 0; i < row_name.length; i++) {
            R row = row_name[i];
            if (!row_to_index.containsKey(row)) {
                row_to_index.put(row, i);
            }
            row_locks[i] = new ReentrantLock(true);
        }

        col_to_index.clear();
        col_Locks = new Lock[col_name.length];
        for (int j = 0; j < col_name.length; j++) {
            C col = col_name[j];
            if (!col_to_index.containsKey(col)) {
                col_to_index.put(col, j);
            }
            col_Locks[j] = new ReentrantLock(true);
        }
    }

    public RealMatrix getData() {
        return data;
    }

    public boolean containRow(R row) {
        return row_to_index.containsKey(row);
    }

    public int getRowIndex(R row) {
        return row_to_index.get(row);
    }

    public boolean containCol(C col) {
        return col_to_index.containsKey(col);
    }

    public int getColIndex(C col) {
        return col_to_index.get(col);
    }

    public void setValue(int i, int j, double v) {
        data.setEntry(i, j, v);
    }

    public void setValue(int i, C col, double v) {
        data.setEntry(i, col_to_index.get(col), v);
    }

    public void setValue(R row, int j, double v) {
        setValue(row_to_index.get(row), j, v);
    }

    public void setValue(R row, C col, double v) {
        setValue(row_to_index.get(row), col_to_index.get(col), v);
    }

    public double getValue(int i, int j) {
        return data.getEntry(i, j);
    }

    public double getValue(int i, C col) {
        return data.getEntry(i, col_to_index.get(col));
    }

    public double getValue(R row, int j) {
        return data.getEntry(row_to_index.get(row), j);
    }

    public double getValue(R row, C col) {
        return data.getEntry(row_to_index.get(row), col_to_index.get(col));
    }

    public R[] getRowName() {
        return row_name;
    }

    public C[] getColName() {
        return col_name;
    }

    public int[] getDimension() {
        return new int[] { row_name.length, col_name.length };
    }

    /**
     * 确保第一个数据框的列的顺序与第二个数据框的行的顺序一致
     * 
     */
    public <T> Matrix<R, T> multi(Matrix<?, T> d) {
        if (this.getColName().length != d.getRowName().length) {
            throw new IllegalArgumentException("dimension not match when matrix multiply");
        }
        Matrix<R, T> new_data = new Matrix<>(this.getRowName(), d.getColName(), 0);
        new_data.data = this.data.multiply(d.data);
        return new_data;
    }

    public void add(int i, int j, double v) {
        data.addToEntry(i, j, v);
    }

    public void rowConcurrentAdd(int i, int j, double v) {
        row_locks[i].lock();
        try {
            add(i, j, v);
        } finally {
            row_locks[i].unlock();
        }
    }

    public void colConcurrentAdd(int i, int j, double v) {
        col_Locks[j].lock();
        try {
            add(i, j, v);
        } finally {
            col_Locks[j].unlock();
        }
    }

    public void add(R row, int j, double v) {
        add(row_to_index.get(row), j, v);
    }

    public void rowConcurrentAdd(R row, int j, double v) {
        rowConcurrentAdd(row_to_index.get(row), j, v);
    }

    public void colConcurrentAdd(R row, int j, double v) {
        colConcurrentAdd(row_to_index.get(row), j, v);
    }

    public void add(R row, C col, double v) {
        add(row_to_index.get(row), col_to_index.get(col), v);
    }

    public void rowConcurrentAdd(R row, C col, double v) {
        rowConcurrentAdd(row_to_index.get(row), col_to_index.get(col), v);
    }

    public void colConcurrentAdd(R row, C col, double v) {
        colConcurrentAdd(row_to_index.get(row), col_to_index.get(col), v);
    }

    public void add(Matrix<R, C> matrix) {
        data = data.add(matrix.data);
    }

    public synchronized void concurrentAdd(Matrix<R, C> matrix) {
        add(matrix);
    }

    public void appendRow(R row_name, double[] value) {
        if (value == null) {
            value = new double[this.col_name.length];
        } else if (value.length != this.col_name.length) {
            value = Arrays.copyOf(value, this.col_name.length);
        }
        R[] new_row = Arrays.copyOf(this.row_name, this.row_name.length + 1);
        new_row[new_row.length - 1] = row_name;
        RealMatrix new_data = MatrixUtils.createRealMatrix(new_row.length, this.col_name.length);
        for (int i = 0; i < this.row_name.length; i++) {
            new_data.setRow(i, this.data.getRow(i));
        }
        new_data.setRow(new_data.getRowDimension() - 1, value);
        this.row_name = new_row;
        this.data = new_data;
        // if (!row_to_index.containsKey(row_name)) {
        // row_to_index.put(row_name, this.row_name.length - 1);
        // }
        this.indexInit();
    }

    public boolean deleteRow(R row_name) {
        if (this.containRow(row_name)) {
            return deleteRow(this.getRowIndex(row_name));
        } else {
            return false;
        }
    }

    public boolean deleteRow(int i) {
        if (i >= this.getRowName().length) {
            return false;
        }
        R[] new_row = Arrays.copyOf(this.row_name, this.row_name.length - 1);
        if (i < this.getRowName().length - 1) {
            System.arraycopy(this.row_name, i + 1, new_row, i, new_row.length - i);
        }
        int[] selected_row_index = new int[this.getRowName().length - 1];
        int[] selected_col_index = new int[this.getColName().length];
        for (int p = 0; p < i; p++) {
            selected_row_index[p] = p;
        }
        for (int p = i; p < selected_row_index.length; p++) {
            selected_row_index[p] = p + 1;
        }
        for (int p = 0; p < selected_col_index.length; p++) {
            selected_col_index[p] = p;
        }
        this.data = this.data.getSubMatrix(selected_row_index, selected_col_index);
        this.row_name = new_row;
        this.indexInit();
        return true;
    }

    public void appendCol(C col_name, double[] value) {
        if (value == null) {
            value = new double[this.row_name.length];
        } else if (value.length != this.row_name.length) {
            value = Arrays.copyOf(value, this.row_name.length);
        }
        C[] new_col = Arrays.copyOf(this.col_name, this.col_name.length + 1);
        new_col[new_col.length - 1] = col_name;
        RealMatrix new_data = MatrixUtils.createRealMatrix(this.row_name.length, new_col.length);
        for (int j = 0; j < this.col_name.length; j++) {
            new_data.setColumn(j, this.data.getColumn(j));
        }
        new_data.setColumn(new_data.getColumnDimension() - 1, value);
        this.col_name = new_col;
        this.data = new_data;
        // if (!col_to_index.containsKey(col_name)) {
        // col_to_index.put(col_name, this.col_name.length - 1);
        // }
        this.indexInit();
    }

    public boolean deleteCol(C col_name) {
        if (this.containCol(col_name)) {
            return deleteCol(this.getColIndex(col_name));
        } else {
            return false;
        }
    }

    public boolean deleteCol(int j) {
        if (j >= this.getColName().length) {
            return false;
        }
        C[] new_col = Arrays.copyOf(this.col_name, this.col_name.length - 1);
        if (j < this.getColName().length - 1) {
            System.arraycopy(this.col_name, j + 1, new_col, j, new_col.length - j);
        }
        int[] selected_row_index = new int[this.getRowName().length];
        int[] selected_col_index = new int[this.getColName().length - 1];
        for (int p = 0; p < j; p++) {
            selected_col_index[p] = p;
        }
        for (int p = j; p < selected_col_index.length; p++) {
            selected_col_index[p] = p + 1;
        }
        for (int p = 0; p < selected_row_index.length; p++) {
            selected_row_index[p] = p;
        }
        this.data = this.data.getSubMatrix(selected_row_index, selected_col_index);
        this.col_name = new_col;
        this.indexInit();
        return true;
    }

    public Matrix<C, R> T() {
        RealMatrix new_data = this.data.transpose();
        return new Matrix<>(col_name, row_name, new_data.getData());
    }

    public void setRow(R row_name, double[] v) {
        setRow(row_to_index.get(row_name), v);
    }

    public void setRow(R row_name, double v) {
        setRow(row_to_index.get(row_name), v);
    }

    public void setRow(int i, double v) {
        for (int j = 0; j < col_name.length; j++) {
            setValue(i, j, v);
        }
    }

    public void setRow(int i, double[] v) {
        for (int j = 0; j < col_name.length; j++) {
            if (j >= v.length) {
                break;
            }
            setValue(i, j, v[j]);
        }
    }

    public void setCol(C col_name, double[] v) {
        setCol(col_to_index.get(col_name), v);
    }

    public void setCol(C col_name, double v) {
        setCol(col_to_index.get(col_name), v);
    }

    public void setCol(int j, double v) {
        for (int i = 0; i < row_name.length; i++) {
            setValue(i, j, v);
        }
    }

    public void setCol(int j, double[] v) {
        for (int i = 0; i < row_name.length; i++) {
            if (i >= v.length) {
                break;
            }
            setValue(i, j, v[i]);
        }
    }

    public double[] getRowData(R row_name) {
        return getRowData(getRowIndex(row_name));
    }

    public double[] getRowData(int row_index) {
        return this.data.getRow(row_index);
    }

    public double getRowSum(int i) {
        double sum = 0;
        for (int j = 0; j < col_name.length; j++) {
            sum += getValue(i, j);
        }
        return sum;
    }

    public double getRowSum(R row_name) {
        return getRowSum(row_to_index.get(row_name));
    }

    public double[] getRowSumArrays() {
        double[] sum = new double[row_name.length];
        for (int i = 0; i < sum.length; i++) {
            sum[i] = getRowSum(i);
        }
        return sum;
    }

    public double[] getColumnData(C col_name) {
        return getColumnData(getColIndex(col_name));
    }

    public double[] getColumnData(int col_index) {
        return this.data.getColumn(col_index);
    }

    public double getColumnSum(int j) {
        double sum = 0;
        for (int i = 0; i < row_name.length; i++) {
            sum += getValue(i, j);
        }
        return sum;
    }

    public double getColumnSum(C col_name) {
        return getColumnSum(col_to_index.get(col_name));
    }

    public double[] getColumnSumArrays() {
        double[] sum = new double[col_name.length];
        for (int j = 0; j < sum.length; j++) {
            sum[j] = getColumnSum(j);
        }
        return sum;
    }

    public double getSum() {
        return Arrays.stream(getRowSumArrays()).sum();
    }

    @Override
    public String toString() {
        StringBuilder str = new StringBuilder();
        str.append(this.name);
        for (C c : col_name) {
            str.append("\t" + c);
        }
        str.append("\n");
        for (int i = 0; i < row_name.length; i++) {
            str.append(row_name[i]);
            for (int j = 0; j < col_name.length; j++) {
                str.append("\t" + getValue(i, j));
            }
            str.append("\n");
        }

        str.setLength(str.length() - 1);

        return str.toString();
    }

    public void toCSV(Writer w) throws IOException {
        w.write(name + "," + StringUtil.join(",", col_name) + "\n");
        for (int i = 0; i < row_name.length; i++) {
            w.write(row_name[i] + ","
                    + StringUtil.join(",", Arrays.stream(getRowData(i)).mapToObj(Double::valueOf).toArray()) + "\n");
        }
    }
}
